package auth

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"net/http"
	"sort"
	"strings"
)

// TODO: 完整实现，基于规划片段

type Signer struct {
	keys map[string]string // accessKey -> secretKey
}

func NewSigner() *Signer {
	// 教学用，写死键
	return &Signer{keys: map[string]string{"AKIAIOSFODNN7EXAMPLE": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}}
}

// 定义辅助结构体和函数
type Credential struct {
	AccessKey string
	Date      string
	Region    string
	Service   string
	Scope     string
}

func parseAuthHeader(auth string) (Credential, []string, string, error) {
	if !strings.HasPrefix(auth, "AWS4-HMAC-SHA256 ") {
		return Credential{}, nil, "", fmt.Errorf("invalid authorization algorithm")
	}

	auth = strings.TrimPrefix(auth, "AWS4-HMAC-SHA256 ")
	// The spec allows for comma or comma-space, but boto3 uses ", ".
	pairs := strings.Split(auth, ", ")

	var credVal, signedHeadersVal, signatureVal string
	for _, pair := range pairs {
		kv := strings.SplitN(pair, "=", 2)
		if len(kv) != 2 {
			return Credential{}, nil, "", fmt.Errorf("invalid auth header pair: %s", pair)
		}
		key, value := kv[0], kv[1]

		switch key {
		case "Credential":
			credVal = value
		case "SignedHeaders":
			signedHeadersVal = value
		case "Signature":
			signatureVal = value
		}
	}

	if credVal == "" || signedHeadersVal == "" || signatureVal == "" {
		return Credential{}, nil, "", fmt.Errorf("missing required fields in Authorization header (Credential, SignedHeaders, Signature)")
	}

	// Parse Credential field
	credParts := strings.Split(credVal, "/")
	if len(credParts) != 5 {
		return Credential{}, nil, "", fmt.Errorf("invalid credential format: expected 5 parts, got %d", len(credParts))
	}
	cred := Credential{
		AccessKey: credParts[0],
		Date:      credParts[1],
		Region:    credParts[2],
		Service:   credParts[3],
		Scope:     fmt.Sprintf("%s/%s/%s/%s", credParts[1], credParts[2], credParts[3], credParts[4]),
	}

	signedHeaders := strings.Split(signedHeadersVal, ";")

	return cred, signedHeaders, signatureVal, nil
}

// canonicalRequest creates the canonical request string for Signature Version 4.
// This is a simplified implementation for educational purposes. A full implementation
// needs to handle many more edge cases according to AWS documentation.
func canonicalRequest(r *http.Request, signedHeaders []string) (string, error) {
	// 1. HTTP Method
	method := r.Method

	// 2. Canonical URI
	uri := r.URL.Path

	// 3. Canonical Query String
	// Query parameters must be sorted alphabetically by key
	query := r.URL.RawQuery
	if query != "" {
		// Parse and sort query parameters
		params := r.URL.Query()
		var keys []string
		for k := range params {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		var parts []string
		for _, k := range keys {
			values := params[k]
			for _, v := range values {
				parts = append(parts, k+"="+v)
			}
		}
		query = strings.Join(parts, "&")
	}

	// 4. Canonical Headers
	// Headers must be lowercased, sorted, and have their values trimmed.
	var headerLines []string
	headers := make(map[string][]string)
	for _, h := range signedHeaders {
		var vals []string
		var ok bool
		// The 'Host' header is special. Go puts it in r.Host instead of r.Header.
		if h == "host" {
			vals = []string{r.Host}
			ok = true
		} else {
			vals, ok = r.Header[http.CanonicalHeaderKey(h)]
		}

		if !ok {
			// According to the spec, if a signed header is not present, the request is invalid.
			// However, for simplicity in our teaching example, we might ignore this.
			// For a stricter implementation, return an error here.
			continue
		}
		// Trim whitespace from header values
		var trimmedVals []string
		for _, v := range vals {
			trimmedVals = append(trimmedVals, strings.TrimSpace(v))
		}
		headers[h] = trimmedVals
	}

	sort.Strings(signedHeaders)
	for _, h := range signedHeaders {
		headerLines = append(headerLines, h+":"+strings.Join(headers[h], ","))
	}
	canonicalHeaders := strings.Join(headerLines, "\n") + "\n"

	// 5. Signed Headers
	signedHeadersStr := strings.Join(signedHeaders, ";")

	// 6. Hashed Payload
	// For GET requests or unsigned payloads, this is a pre-calculated string.
	// For PUT/POST, you'd calculate the SHA256 hash of the request body.
	hashedPayload := r.Header.Get("X-Amz-Content-Sha256")
	if hashedPayload == "" {
		// This is common for pre-signed URLs or streaming uploads
		hashedPayload = "UNSIGNED-PAYLOAD"
	}

	return fmt.Sprintf("%s\n%s\n%s\n%s\n%s\n%s", method, uri, query, canonicalHeaders, signedHeadersStr, hashedPayload), nil
}

func stringToSign(amzDate, credentialScope, canonicalRequest string) string {
	hasher := sha256.New()
	hasher.Write([]byte(canonicalRequest))
	hashedCanonicalRequest := hex.EncodeToString(hasher.Sum(nil))

	return fmt.Sprintf("AWS4-HMAC-SHA256\n%s\n%s\n%s", amzDate, credentialScope, hashedCanonicalRequest)
}

func deriveSigningKey(secret, date, region, service string) []byte {
	kDate := hmacSHA256([]byte("AWS4"+secret), []byte(date))
	kRegion := hmacSHA256(kDate, []byte(region))
	kService := hmacSHA256(kRegion, []byte(service))
	return hmacSHA256(kService, []byte("aws4_request"))
}

func hmacSHA256(key, data []byte) []byte {
	h := hmac.New(sha256.New, key)
	h.Write(data)
	return h.Sum(nil)
}

type ctxUserKey string

var ctxUser = ctxUserKey("S3User")

func (s *Signer) Verify(r *http.Request) error {
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" {
		return fmt.Errorf("missing Authorization header")
	}

	cred, signedHeaders, signature, err := parseAuthHeader(authHeader)
	if err != nil {
		return fmt.Errorf("failed to parse Authorization header: %w", err)
	}

	canonical, err := canonicalRequest(r, signedHeaders)
	if err != nil {
		return fmt.Errorf("failed to create canonical request: %w", err)
	}

	amzDate := r.Header.Get("X-Amz-Date")
	stringToSign := stringToSign(amzDate, cred.Scope, canonical)

	secret := s.keys[cred.AccessKey]
	if secret == "" {
		return fmt.Errorf("unknown access key")
	}

	derivedKey := deriveSigningKey(secret, cred.Date, cred.Region, cred.Service)
	expectedSignature := hex.EncodeToString(hmacSHA256(derivedKey, []byte(stringToSign)))

	if subtle.ConstantTimeCompare([]byte(expectedSignature), []byte(signature)) != 1 {
		return fmt.Errorf("signature mismatch")
	}

	ctx := context.WithValue(r.Context(), ctxUser, cred.AccessKey)
	*r = *r.WithContext(ctx)
	return nil
}

// Note: framework-specific middlewares are implemented in server layers (gin/stdhttp).

// Exported thin wrappers to aid testing and external signing
// These expose the simplified SigV4 helpers for building signed requests in tests.

// CanonicalRequest exports canonicalRequest.
func CanonicalRequest(r *http.Request, signedHeaders []string) (string, error) {
	return canonicalRequest(r, signedHeaders)
}

// StringToSign exports stringToSign.
func StringToSign(amzDate, credentialScope, canonical string) string {
	return stringToSign(amzDate, credentialScope, canonical)
}

// DeriveSigningKey exports deriveSigningKey.
func DeriveSigningKey(secret, date, region, service string) []byte {
	return deriveSigningKey(secret, date, region, service)
}

// HexHMAC returns hex-encoded HMAC-SHA256.
func HexHMAC(key, data []byte) string { return hex.EncodeToString(hmacSHA256(key, data)) }
